/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2022-2022. All rights reserved.
 * Description: alg remote lookup.
 * Author: qingjicheng
 * Create: 2022-11-28
 */

#include "alg_remote_lookup.h"
#include "log.h"
#include "dtype_common.h"
#include "adapter_hccp.h"

using namespace std;

namespace hccl {
namespace {
    constexpr u32 LOOKUP_INFO_DATA_COUNT = 1;
}
constexpr u32 MULTIPLES = 2;

AlgRemoteLookUp::AlgRemoteLookUp() : funcTable_({
    {EmbeddingServiceParam::OperatorType::REMOTE_LOOK_UP,
        std::bind(&AlgRemoteLookUp::AlgRemoteLookup, this, std::placeholders::_1)},
    {EmbeddingServiceParam::OperatorType::COLL_REMOTE_LOOK_UP,
        std::bind(&AlgRemoteLookUp::AlgCollRemoteLookup, this, std::placeholders::_1)},
    {EmbeddingServiceParam::OperatorType::COLL_REMOTE_LOOK_UP_PAIRED,
        std::bind(&AlgRemoteLookUp::AlgCollRemoteLookupPaired, this, std::placeholders::_1)},
    {EmbeddingServiceParam::OperatorType::GET_LOOK_UP_REQUEST,
        std::bind(&AlgRemoteLookUp::AlgGetLookupRequest, this, std::placeholders::_1)},
    {EmbeddingServiceParam::OperatorType::WAIT_LOOK_UP_DATA,
        std::bind(&AlgRemoteLookUp::AlgWaitLookupData, this, std::placeholders::_1)},
    {EmbeddingServiceParam::OperatorType::SET_LOOK_UP_RESPONSE,
        std::bind(&AlgRemoteLookUp::AlgSetLookupResponse, this, std::placeholders::_1)},
    {EmbeddingServiceParam::OperatorType::WAIT_SOME,
        std::bind(&AlgRemoteLookUp::AlgWaitSome, this, std::placeholders::_1)},
    {EmbeddingServiceParam::OperatorType::SERVICE_CANCEL,
        std::bind(&AlgRemoteLookUp::AlgServiceCancel, this, std::placeholders::_1)},
    {EmbeddingServiceParam::OperatorType::REMOTE_LOOKUP_KEYS_DUPLICATES,
        std::bind(&AlgRemoteLookUp::AlgKeysDuplicates, this, std::placeholders::_1)},
    {EmbeddingServiceParam::OperatorType::REMOTE_LOOKUP_SEND_KEYS,
        std::bind(&AlgRemoteLookUp::AlgSendKeys, this, std::placeholders::_1)},
    {EmbeddingServiceParam::OperatorType::REMOTE_LOOKUP_RECV_VALUES,
        std::bind(&AlgRemoteLookUp::AlgRecvValues, this, std::placeholders::_1)},
    {EmbeddingServiceParam::OperatorType::REMOTE_LOOKUP_RECOVER_VALUE,
        std::bind(&AlgRemoteLookUp::AlgRecoverValue, this, std::placeholders::_1)},
    {EmbeddingServiceParam::OperatorType::REMOTE_LOOKUP_RESET_UNIQUE_HANDLE,
        std::bind(&AlgRemoteLookUp::AlgResetUniqueHandle, this, std::placeholders::_1)},
    {EmbeddingServiceParam::OperatorType::REMOTE_LOOKUP_RESET_UNIQUE_TASK,
        std::bind(&AlgRemoteLookUp::AlgResetUniqueTask, this, std::placeholders::_1)},
    {EmbeddingServiceParam::OperatorType::REMOTE_LOOKUP_WAIT_SEND_KEY_FINISH,
        std::bind(&AlgRemoteLookUp::AlgWaitSendKeyFinish, this, std::placeholders::_1)}
    })
{
}

HcclResult AlgRemoteLookUp::Process(HcclOpPtr opPtr)
{
    CHK_PTR_NULL(opPtr);
    auto op = static_cast<HcclOperator<EmbeddingServiceParam> *>(opPtr);
    CHK_PRT_RET(funcTable_.find(op->param_.opType) == funcTable_.end(),
        HCCL_ERROR("[Process]funcTable_ without opType[%u]", op->param_.opType), HCCL_E_PARA);
    CHK_RET(funcTable_[op->param_.opType](op));
    return HCCL_SUCCESS;
}

HcclResult AlgRemoteLookUp::AlgRemoteLookup(HcclOperator<EmbeddingServiceParam> *opPtr)
{
    PartitionMapTask *partitionMapTaskPtr = new (std::nothrow) PartitionMapTask();
    CHK_PTR_NULL(partitionMapTaskPtr);
    EmbeddingServiceParam &param = opPtr->param_;
    *partitionMapTaskPtr = param;
    partitionMapTaskPtr->taskType = OpTaskType::DROP_DUPLICATES;
    partitionMapTaskPtr->waitLookUpRanksPtr = &opPtr->param_.embeddingParam.waitLookUpRanks;

    OpTaskPtr taskPtr(partitionMapTaskPtr);
    opPtr->PushTask(std::move(taskPtr));
    return HCCL_SUCCESS;
}

void AlgRemoteLookUp::AlgCollRemoteLookupParaAssign(const HcclOperator<EmbeddingServiceParam> &opPtr,
    PartitionMapTask &partitionMapTask)
{
    const EmbeddingServiceParam &param = opPtr.param_;
    partitionMapTask = param;
    u32 valueItemSize = opPtr.param_.embeddingParam.valueItemSize;
    partitionMapTask.taskType = OpTaskType::REMOTE_LOOKUP_PROCESS;

    partitionMapTask.keyTransferMemSize = partitionMapTask.keyMaxNum * sizeof(s64);
    partitionMapTask.valueTransferMemSize = partitionMapTask.keyMaxNum * valueItemSize;

    partitionMapTask.enableKeyCounter = (opPtr.param_.embeddingParam.flags == ES_FLAGS_ENABLE_COUNTER) ? true : false;
    HCCL_DEBUG("CollRemoteLookup enableKeyCounter[%d], disableUnique[%d], intZerocpyFlag[%d], outZerocpyFlag[%d]",
        partitionMapTask.enableKeyCounter, partitionMapTask.disableUnique, partitionMapTask.intZerocpyFlag,
        partitionMapTask.outZerocpyFlag);
}

HcclResult AlgRemoteLookUp::AlgCollRemoteLookup(HcclOperator<EmbeddingServiceParam> *opPtr)
{
    u64 keyMaxNum = opPtr->param_.keyMaxNum;
    // keyMaxNum大于ES_PIPELINE_THRESHOLD的部分进行流水。
    bool usePipeline = (keyMaxNum > ES_PIPELINE_THRESHOLD);

    u32 cubeIdx = 0;
    for (u64 keyNum = 0; keyNum < keyMaxNum; keyNum += ES_PIPELINE_THRESHOLD) {
        PartitionMapTask *partitionMapTaskPtr = new (nothrow) PartitionMapTask();
        CHK_PTR_NULL(partitionMapTaskPtr);

        AlgCollRemoteLookupParaAssign(*opPtr, *partitionMapTaskPtr);

        partitionMapTaskPtr->cubeIndex = cubeIdx;
        partitionMapTaskPtr->usePipeline = usePipeline;
        cubeIdx++;

        OpTaskPtr taskPtr(partitionMapTaskPtr);
        opPtr->PushTask(std::move(taskPtr));
    }

    if (!usePipeline) {
        return HCCL_SUCCESS;
    }

    // reset去重句柄， TaskResetUniqueHandleKernelLaunch
    ResetUniqueHandleTask *resetUniqueHandleTaskPtr = new (nothrow) ResetUniqueHandleTask();
    CHK_PTR_NULL(resetUniqueHandleTaskPtr);

    resetUniqueHandleTaskPtr->tag = opPtr->param_.tag;
    resetUniqueHandleTaskPtr->taskType = OpTaskType::RESET_UNIQUE_HANDLE;
    resetUniqueHandleTaskPtr->stream = opPtr->param_.embeddingParam.stream;
    resetUniqueHandleTaskPtr->subStreamInfo = opPtr->param_.embeddingParam.subStreamInfo;
    resetUniqueHandleTaskPtr->notifyInfo = opPtr->param_.embeddingParam.notifyInfo;
    resetUniqueHandleTaskPtr->usePipeline = usePipeline;
    OpTaskPtr taskPtr1(resetUniqueHandleTaskPtr);
    opPtr->PushTask(std::move(taskPtr1));

    // wait send keys finish
    ResetUniqueHandleTask *waitSendKeyFinishTaskPtr = new (nothrow) ResetUniqueHandleTask();
    CHK_PTR_NULL(waitSendKeyFinishTaskPtr);

    waitSendKeyFinishTaskPtr->tag = opPtr->param_.tag;
    waitSendKeyFinishTaskPtr->taskType = OpTaskType::WAIT_SEND_KEY_FINISH;
    waitSendKeyFinishTaskPtr->stream = opPtr->param_.embeddingParam.stream;
    waitSendKeyFinishTaskPtr->subStreamInfo = opPtr->param_.embeddingParam.subStreamInfo;
    waitSendKeyFinishTaskPtr->notifyInfo = opPtr->param_.embeddingParam.notifyInfo;
    waitSendKeyFinishTaskPtr->usePipeline = usePipeline;
    OpTaskPtr taskPtr2(waitSendKeyFinishTaskPtr);
    opPtr->PushTask(std::move(taskPtr2));

    return HCCL_SUCCESS;
}

HcclResult AlgRemoteLookUp::AlgCollRemoteLookupPaired(HcclOperator<EmbeddingServiceParam> *opPtr)
{
    return AlgCollRemoteLookup(opPtr);
}

HcclResult AlgRemoteLookUp::AlgKeysDuplicates(HcclOperator<EmbeddingServiceParam> *opPtr)
{
    PartitionMapTask *partitionMapTaskPtr = new (std::nothrow) PartitionMapTask();
    CHK_PTR_NULL(partitionMapTaskPtr);
    EmbeddingServiceParam &param = opPtr->param_;
    *partitionMapTaskPtr = param;
    HCCL_INFO("keyMaxNum is %u.", opPtr->param_.keyMaxNum);
    partitionMapTaskPtr->taskType = OpTaskType::KEYS_DUPLICATES;

    OpTaskPtr taskPtr(partitionMapTaskPtr);
    opPtr->PushTask(std::move(taskPtr));
    return HCCL_SUCCESS;
}

HcclResult AlgRemoteLookUp::AlgSendKeys(HcclOperator<EmbeddingServiceParam> *opPtr)
{
    PartitionMapTask *partitionMapTaskPtr = new (std::nothrow) PartitionMapTask();
    CHK_PTR_NULL(partitionMapTaskPtr);
    EmbeddingServiceParam &param = opPtr->param_;
    *partitionMapTaskPtr = param;
    partitionMapTaskPtr->taskType = OpTaskType::SEND;

    HCCL_DEBUG("AlgSendKeys partitionMapTaskPtr->rdmaEnveInfosTransferMem[%llu], haveRdmaConn[%d]",
        hash<void *>{}(partitionMapTaskPtr->rdmaEnveInfosTransferMem), partitionMapTaskPtr->haveRdmaConn);

    OpTaskPtr taskPtr(partitionMapTaskPtr);
    opPtr->PushTask(std::move(taskPtr));
    return HCCL_SUCCESS;
}

HcclResult AlgRemoteLookUp::AlgRecvValues(HcclOperator<EmbeddingServiceParam> *opPtr)
{
    PartitionMapTask *partitionMapTaskPtr = new (std::nothrow) PartitionMapTask();
    CHK_PTR_NULL(partitionMapTaskPtr);
    EmbeddingServiceParam &param = opPtr->param_;
    *partitionMapTaskPtr = param;
    partitionMapTaskPtr->taskType = OpTaskType::RECV;

    OpTaskPtr taskPtr(partitionMapTaskPtr);
    opPtr->PushTask(std::move(taskPtr));
    return HCCL_SUCCESS;
}

HcclResult AlgRemoteLookUp::AlgRecoverValue(HcclOperator<EmbeddingServiceParam> *opPtr)
{
    PartitionMapTask *partitionMapTaskPtr = new (std::nothrow) PartitionMapTask();
    CHK_PTR_NULL(partitionMapTaskPtr);
    EmbeddingServiceParam &param = opPtr->param_;
    *partitionMapTaskPtr = param;
    partitionMapTaskPtr->taskType = OpTaskType::RECOVER_VALUE;

    OpTaskPtr taskPtr(partitionMapTaskPtr);
    opPtr->PushTask(std::move(taskPtr));
    return HCCL_SUCCESS;
}

HcclResult AlgRemoteLookUp::AlgResetUniqueTask(HcclOperator<EmbeddingServiceParam> *opPtr)
{
    PartitionMapTask *partitionMapTaskPtr = new (std::nothrow) PartitionMapTask();
    CHK_PTR_NULL(partitionMapTaskPtr);
    EmbeddingServiceParam &param = opPtr->param_;
    *partitionMapTaskPtr = param;
    partitionMapTaskPtr->taskType = OpTaskType::RESET_UNIQUE_TASK;

    OpTaskPtr taskPtr(partitionMapTaskPtr);
    opPtr->PushTask(std::move(taskPtr));
    return HCCL_SUCCESS;
}

HcclResult AlgRemoteLookUp::AlgResetUniqueHandle(HcclOperator<EmbeddingServiceParam> *opPtr)
{
    PartitionMapTask *partitionMapTaskPtr = new (std::nothrow) PartitionMapTask();
    CHK_PTR_NULL(partitionMapTaskPtr);
    partitionMapTaskPtr->taskType = OpTaskType::RESET_UNIQUE_HANDLE_KERNEL_LAUNCH;

    OpTaskPtr taskPtr(partitionMapTaskPtr);
    opPtr->PushTask(std::move(taskPtr));
    return HCCL_SUCCESS;
}

HcclResult AlgRemoteLookUp::AlgWaitSendKeyFinish(HcclOperator<EmbeddingServiceParam> *opPtr)
{
    PartitionMapTask *partitionMapTaskPtr = new (std::nothrow) PartitionMapTask();
    CHK_PTR_NULL(partitionMapTaskPtr);
    partitionMapTaskPtr->taskType = OpTaskType::WAIT_SEND_KEY_FINISH;

    OpTaskPtr taskPtr(partitionMapTaskPtr);
    opPtr->PushTask(std::move(taskPtr));
    return HCCL_SUCCESS;
}

HcclResult AlgRemoteLookUp::AlgGetLookupRequest(HcclOperator<EmbeddingServiceParam> *opPtr)
{
    TransTask *transTaskPtr = new (std::nothrow) TransTask();
    CHK_PTR_NULL(transTaskPtr);
    transTaskPtr->tag = opPtr->param_.tag;
    transTaskPtr->count = LOOKUP_INFO_DATA_COUNT;
    transTaskPtr->improbeRanks = opPtr->param_.getRequestParam.workerList;
    transTaskPtr->recvMem = opPtr->param_.keys;
    transTaskPtr->dataType = HCCL_DATA_TYPE_UINT64;
    transTaskPtr->outputs.push_back(opPtr->param_.getRequestParam.reqStatus);
    transTaskPtr->taskType = OpTaskType::RECV;
    transTaskPtr->srcRank = opPtr->param_.userRank;
    OpTaskPtr taskPtr(transTaskPtr);
    opPtr->PushTask(std::move(taskPtr));
    return HCCL_SUCCESS;
}

HcclResult AlgRemoteLookUp::AlgWaitLookupData(HcclOperator<EmbeddingServiceParam> *opPtr)
{
    u64 offset = 0;
    for (auto &pair : opPtr->param_.embeddingParam.waitLookUpRanks) {
        TransTask *transTaskPtr = new (std::nothrow) TransTask();
        CHK_PTR_NULL(transTaskPtr);
        transTaskPtr->count = pair.second;
        transTaskPtr->dataType = HCCL_DATA_TYPE_INT8; // embedding表数据为fp32
        transTaskPtr->recvMem = static_cast<u8 *>(opPtr->param_.values)
            + offset;
        transTaskPtr->dstRank = pair.first;
        transTaskPtr->stream = opPtr->param_.embeddingParam.stream;
        transTaskPtr->taskType = OpTaskType::RECV;
        transTaskPtr->outZerocpyFlag = opPtr->param_.outZerocpyFlag;
        transTaskPtr->srcRank = opPtr->param_.userRank;
        transTaskPtr->tag = opPtr->param_.tag;
        OpTaskPtr taskPtr(transTaskPtr);
        opPtr->PushTask(std::move(taskPtr));
        offset += pair.second;
    }
    return HCCL_SUCCESS;
}

HcclResult AlgRemoteLookUp::AlgSetLookupResponse(HcclOperator<EmbeddingServiceParam> *opPtr)
{
    void *sendMem = nullptr;
    if (!opPtr->param_.getRequestParam.transferMems.empty()) {
        OpTaskPtr taskPtr(new (std::nothrow) MemTask());
        CHK_SMART_PTR_NULL(taskPtr);
        MemTask *memTaskPtr = static_cast<MemTask *>(taskPtr.get());
        memTaskPtr->dstMem = opPtr->param_.getRequestParam.transferMems.begin()->second[1].ptr();
        memTaskPtr->maxSize = opPtr->param_.getRequestParam.transferMems.begin()->second[1].size();
        memTaskPtr->srcMem = opPtr->param_.values;
        memTaskPtr->size = opPtr->param_.valueCount * SIZE_TABLE[opPtr->param_.valuesType];
        memTaskPtr->taskType = OpTaskType::MEMCPY;
        sendMem = memTaskPtr->dstMem;
        opPtr->PushTask(std::move(taskPtr));
    }

    TransTask *transTaskPtr = new (std::nothrow) TransTask();
    CHK_PTR_NULL(transTaskPtr);
    transTaskPtr->count = opPtr->param_.valueCount;
    transTaskPtr->dataType = opPtr->param_.valuesType;
    transTaskPtr->sendMem = sendMem == nullptr ? opPtr->param_.values : sendMem;
    transTaskPtr->improbeRanks = opPtr->param_.getRequestParam.workerList;
    transTaskPtr->outRequest = opPtr->param_.getRequestParam.outRequest;
    transTaskPtr->taskType = OpTaskType::ISEND;
    transTaskPtr->tag = opPtr->param_.tag;
    transTaskPtr->srcRank = opPtr->param_.userRank;
    OpTaskPtr taskPtr(transTaskPtr);
    opPtr->PushTask(std::move(taskPtr));
    return HCCL_SUCCESS;
}

HcclResult AlgRemoteLookUp::AlgWaitSome(HcclOperator<EmbeddingServiceParam> *opPtr)
{
    WaitSomeTask *transTaskPtr = new (std::nothrow) WaitSomeTask();
    CHK_PTR_NULL(transTaskPtr);
    transTaskPtr->requestCount = opPtr->param_.waitSomeParam.requestCount;
    transTaskPtr->requestArray = opPtr->param_.waitSomeParam.requestArray;
    transTaskPtr->compCount = opPtr->param_.waitSomeParam.compCount;
    transTaskPtr->compIndices = opPtr->param_.waitSomeParam.compIndices;
    transTaskPtr->compStatus = opPtr->param_.waitSomeParam.compStatus;
    transTaskPtr->taskType = OpTaskType::WAIT;
    OpTaskPtr taskPtr(transTaskPtr);
    opPtr->PushTask(std::move(taskPtr));
    return HCCL_SUCCESS;
}
} // namespace hccl
